#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
完整的MES系统Django兼容性修复脚本
"""

import os
import re
import sys

def fix_foreign_key_fields(root_dir):
    """修复所有ForeignKey字段缺少on_delete参数的问题"""
    
    # 模式匹配各种ForeignKey用法
    patterns = [
        # 基本的ForeignKey定义
        (r'(\w+)\s*=\s*models\.ForeignKey\(([^,)]+)(?:,([^)]*))?\)', 
         r'\1 = models.ForeignKey(\2, on_delete=models.CASCADE\3)'),
        
        # TreeForeignKey定义
        (r'(\w+)\s*=\s*TreeForeignKey\(([^,)]+)(?:,([^)]*))?\)', 
         r'\1 = TreeForeignKey(\2, on_delete=models.CASCADE\3)'),
        
        # 其他ForeignKey变体
        (r'(\w+)\s*=\s*ForeignKey\(([^,)]+)(?:,([^)]*))?\)', 
         r'\1 = ForeignKey(\2, on_delete=models.CASCADE\3)'),
    ]
    
    fixed_files = []
    
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for filename in filenames:
            if filename.endswith('.py'):
                filepath = os.path.join(dirpath, filename)
                
                try:
                    with open(filepath, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    original_content = content
                    
                    # 跳过已经有on_delete参数的字段
                    if 'on_delete=' in content:
                        # 只修复没有on_delete的字段
                        for pattern, replacement in patterns:
                            # 只替换没有on_delete参数的ForeignKey
                            matches = re.findall(pattern, content)
                            for match in matches:
                                full_match = match[0] + ' = ' + ('models.ForeignKey(' if 'models.ForeignKey' in pattern else 'TreeForeignKey(' if 'TreeForeignKey' in pattern else 'ForeignKey(') + match[1]
                                if len(match) > 2 and match[2]:
                                    full_match += ', ' + match[2]
                                full_match += ')'
                                
                                # 检查这个具体的字段是否已经有on_delete参数
                                if 'on_delete=' not in full_match:
                                    content = re.sub(pattern, replacement, content, count=1)
                    else:
                        # 如果文件中没有on_delete，则进行全面替换
                        for pattern, replacement in patterns:
                            content = re.sub(pattern, replacement, content)
                    
                    # 如果内容有变化，写回文件
                    if content != original_content:
                        with open(filepath, 'w', encoding='utf-8') as f:
                            f.write(content)
                        fixed_files.append(filepath)
                        print(f"Fixed ForeignKey in: {filepath}")
                        
                except Exception as e:
                    print(f"Error processing {filepath}: {e}")
    
    return fixed_files

def fix_mptt_compatibility(root_dir):
    """修复django-mptt兼容性问题"""
    
    fixed_files = []
    
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for filename in filenames:
            if filename.endswith('.py'):
                filepath = os.path.join(dirpath, filename)
                
                try:
                    with open(filepath, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    original_content = content
                    
                    # 修复mptt导入
                    content = re.sub(
                        r'from mptt.models import MPTTModel,TreeForeignKey',
                        'from mptt.models import MPTTModel, TreeForeignKey',
                        content
                    )
                    
                    # 修复mptt字段
                    content = re.sub(
                        r'TreeForeignKey\(([^,)]+),\s*([^)]*)\)',
                        r'TreeForeignKey(\1, on_delete=models.CASCADE, \2)',
                        content
                    )
                    
                    # 清理重复的on_delete参数
                    content = re.sub(
                        r'on_delete=models\.CASCADE,\s*on_delete=models\.CASCADE',
                        'on_delete=models.CASCADE',
                        content
                    )
                    
                    if content != original_content:
                        with open(filepath, 'w', encoding='utf-8') as f:
                            f.write(content)
                        fixed_files.append(filepath)
                        print(f"Fixed MPTT in: {filepath}")
                        
                except Exception as e:
                    print(f"Error processing {filepath}: {e}")
    
    return fixed_files

def fix_django_imports(root_dir):
    """修复Django导入兼容性问题"""
    
    import_fixes = [
        # URL导入修复
        (r'from django\.conf\.urls import url', 'from django.urls import re_path as url'),
        (r'from django\.conf\.urls import include', 'from django.urls import include'),
        
        # 翻译导入修复
        (r'from django\.utils\.translation import ugettext_lazy', 'from django.utils.translation import gettext_lazy'),
        (r'from django\.utils\.translation import ugettext', 'from django.utils.translation import gettext'),
        
        # 其他常见导入修复
        (r'from django\.core\.urlresolvers import reverse', 'from django.urls import reverse'),
        (r'from django\.contrib\.auth\.models import User', 'from django.contrib.auth.models import User'),
    ]
    
    fixed_files = []
    
    for dirpath, dirnames, filenames in os.walk(root_dir):
        for filename in filenames:
            if filename.endswith('.py'):
                filepath = os.path.join(dirpath, filename)
                
                try:
                    with open(filepath, 'r', encoding='utf-8') as f:
                        content = f.read()
                    
                    original_content = content
                    
                    for old_import, new_import in import_fixes:
                        content = re.sub(old_import, new_import, content)
                    
                    if content != original_content:
                        with open(filepath, 'w', encoding='utf-8') as f:
                            f.write(content)
                        fixed_files.append(filepath)
                        print(f"Fixed imports in: {filepath}")
                        
                except Exception as e:
                    print(f"Error processing {filepath}: {e}")
    
    return fixed_files

def main():
    """主函数"""
    
    print("=== MES Django兼容性修复脚本 ===")
    
    # 获取apps目录
    apps_dir = os.path.join(os.path.dirname(__file__), 'apps')
    
    if not os.path.exists(apps_dir):
        print(f"错误：未找到apps目录：{apps_dir}")
        return
    
    print(f"开始修复目录：{apps_dir}")
    
    # 1. 修复ForeignKey字段
    print("\n1. 修复ForeignKey字段...")
    foreign_key_files = fix_foreign_key_fields(apps_dir)
    print(f"修复了 {len(foreign_key_files)} 个文件的ForeignKey问题")
    
    # 2. 修复MPTT兼容性
    print("\n2. 修复MPTT兼容性...")
    mptt_files = fix_mptt_compatibility(apps_dir)
    print(f"修复了 {len(mptt_files)} 个文件的MPTT问题")
    
    # 3. 修复Django导入
    print("\n3. 修复Django导入...")
    import_files = fix_django_imports(apps_dir)
    print(f"修复了 {len(import_files)} 个文件的导入问题")
    
    # 统计
    all_fixed_files = set(foreign_key_files + mptt_files + import_files)
    print(f"\n=== 修复完成 ===")
    print(f"总共修复了 {len(all_fixed_files)} 个文件")
    
    return len(all_fixed_files)

if __name__ == '__main__':
    main() 